from flask import jsonify, request, Response
from datetime import datetime, time
from models import Question, Choice, Bank, QuestionBank, Answer, Exercise, Knowledge
from app.extensions import db
import json
import io
from flask import send_file
from docx import Document
from . import problem_bp
from openai import OpenAI
import re
import os
import traceback
import pandas as pd
from langchain.agents import load_tools, AgentType
from langchain.agents import initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain import PromptTemplate
from typing import List, Dict
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.llms import OpenAI


@problem_bp.route('/')
def index():
    return jsonify(u"this is problem url prefix")


@problem_bp.route('/summary')
def exercise_summary():
    exercise_id = request.args.get('exercise_id')
    if not exercise_id:
        return jsonify({'error': 'Exercise ID is required'}), 400

    # Fetch the exercise
    exercise = db.session.query(Exercise).filter_by(exercise_id=exercise_id).first()
    if not exercise:
        return jsonify({'error': 'Exercise not found'}), 404

    # Get the bank associated with the exercise
    bank_id = exercise.exercise_bank

    # Get questions associated with the bank
    questions = db.session.query(Question).join(QuestionBank).filter(QuestionBank.bank_id == bank_id).all()

    # Fetch choices and user answers, determine correctness

    exercise_details = []
    question_cnt = 0
    correct_cnt = 0
    incorrect_cnt = 0
    unanswered_cnt = 0
    correct_rate = 0.0
    df = pd.DataFrame([], [], ['正确', '错误', '未作答'])
    for question in questions:
        choices = db.session.query(Choice).filter_by(question_id=question.question_id).all()
        user_answer = db.session.query(Answer).filter_by(question_id=question.question_id,
                                                         exercise_id=exercise_id).first()

        # Format user choice
        question_cnt = question_cnt + 1
        user_choice = -1
        is_correct = '未作答'
        if user_answer.choice != 0:
            user_choice = 0
            is_correct = '错误'
            for choice in choices:
                if choice.choice_id == user_answer.choice:
                    if choice.is_correct:
                        is_correct = '正确'
                        correct_cnt = correct_cnt + 1
                    else:
                        incorrect_cnt = incorrect_cnt + 1
                    break
                user_choice = user_choice + 1
        else:
            unanswered_cnt = unanswered_cnt + 1

        # Format choices
        formatted_choices = [
            {'choice_id': choice.choice_id, 'choice_text': choice.choice_text, 'is_correct': choice.is_correct} for
            choice in choices]

        # Add question details
        exercise_details.append({
            'knowledge': question.tag,
            'question_id': question.question_id,
            'question_text': question.question_text,
            'choices': formatted_choices,
            'is_correct': is_correct,
            'user_choice': user_choice,
            'description': question.description
        })

        # update graph
        if question.tag not in df.index.values:
            df.loc[question.tag] = [0, 0, 0]

        df.loc[question.tag][is_correct] += 1

    if question_cnt != 0:
        correct_rate = round(correct_cnt / question_cnt * 100, 2)

    category = df.index.values.tolist()  # Avoid error: Object of type ndarray is not JSON serializable
    series = list(df.to_dict(orient='list').values())

    # Return the data
    return jsonify({
        'exercise_details': exercise_details,
        'correct_cnt': correct_cnt,
        'incorrect_cnt': incorrect_cnt,
        'unanswered_cnt': unanswered_cnt,
        'correct_rate': correct_rate,
        'main_knowledge': ";".join(category),
        'category': category,
        'series': series
    })


def rec_construct_topic_tree(knowledge_all, current_node):
    for knowledge in knowledge_all:
        if knowledge.parentid == current_node['id']:
            temp = {'id': knowledge.id, 'value': knowledge.name, 'label': knowledge.name, 'children': []}
            rec_construct_topic_tree(knowledge_all, temp)
            current_node['children'].append(temp)


@problem_bp.route('/topic')
def get_topic():
    knowledge_all = db.session.query(Knowledge).all()

    # 构建json树用来返回
    res = []
    for knowledge in knowledge_all:
        # 根节点（多个）
        if knowledge.parentid == '0':
            temp = {'id': knowledge.id, 'value': knowledge.name, 'label': knowledge.name, 'children': []}
            rec_construct_topic_tree(knowledge_all, temp)
            res.append(temp)

    # Return the data
    return jsonify({'topics': res})


@problem_bp.route('/download', methods=["POST"])
def download_exercise():
    download_list = request.form.getlist('exportList[]')
    print(request.form)
    download_type = request.form.get('suffix')
    print(download_type)
    if download_type == 'docx':
        doc = Document()
        doc.add_heading('练习题', 0)
        indexes = 1
        # 遍历问题列表
        for question_id in download_list:
            # 查询问题
            question = db.session.query(Question).filter(Question.question_id == question_id).first()
            if question:
                doc.add_heading(f"题目{indexes}：{question.question_text}", level=1)
                # 查询当前问题的选项
                choices = db.session.query(Choice).filter(Choice.question_id == question.question_id).all()
                option_letter = 'A'
                for choice in choices:
                    # 添加选项到文档
                    doc.add_paragraph(f" {option_letter}：{choice.choice_text}")
                    option_letter = chr(ord(option_letter) + 1)  # Increment letter to the next one
            indexes += 1

        fake_file = io.BytesIO()
        doc.save(fake_file)
        fake_file.seek(0)
        return send_file(fake_file, mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document')
    elif download_type == 'json':
        download_json(download_list)
    else:
        download_excel(download_excel)
    return jsonify("error"), 500


@problem_bp.route('/generate', methods=["POST"])
def generate_problem():
    """Create a completion."""
    cnt = int(request.form.get('cnt'))
    print(cnt,type(cnt))
    knowledge = request.form.get('knowledge')
    isScene = request.form.get('isScene')
    knowledge = knowledge_format(knowledge)
    question_type = request.form.get('type')
    otherRestriction = request.form.get('otherRestriction')

    if isScene:
        otherRestriction += "你需要为题目编一个完整、有教育意义的故事背景，创设的题目需要以生活中真实发生的问题为背景，使用在生活中会真实出现的数字，问题是要解决生活中实际出现的问题"

    difficulty_list = ['初级', '高级', '梯度增强']
    difficulty_level = int(request.form.get('difficulty'))
    difficulty = difficulty_list[difficulty_level]
    openai_api_key = os.environ["OPENAI_API_KEY"]
    example_problem = search_example_problems(knowledge)
    generated_questions = generate_questions(cnt, knowledge, difficulty, otherRestriction, example_problem)
    question_with_answer = solve_problems(generated_questions)
    question_completed = add_choices(question_with_answer)
    parsed_problems = store_question(question_completed, knowledge, difficulty_level)
    return jsonify(parsed_problems),200


def search_example_problems(knowledge):
    llm = ChatOpenAI(temperature=0.5, model_name="gpt-4-turbo-preview")
    # 谷歌搜索 key
    serpapi_api_key = os.environ["SERPAPI_API_KEY"]
    toolkit = load_tools(["serpapi"], llm=llm, serpapi_api_key=serpapi_api_key)
    agent = initialize_agent(
        tools=toolkit,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True
    )
    template = "{value}，请使用中文回答"
    prompt = PromptTemplate(
        input_variables=["value"],
        template=template,
    )

    questions = f"请帮我找几道中小学{knowledge}的典型应用题，并总结一下常见题型"
    print("例题搜索结果: \n")
    example_problem = agent.run(prompt.format(value=questions))
    print(example_problem)

    return example_problem


class QuestionList(BaseModel):
    question_contents: List[str] = Field(description="list of question contents")


def generate_questions(cnt, knowledge, difficulty, otherRestriction, example_problem):
    question_list = []
    for i in range(cnt):
        llm = ChatOpenAI(temperature=0.5, model_name="gpt-4-turbo-preview")
        # 谷歌搜索 key
        serpapi_api_key = os.environ["SERPAPI_API_KEY"]
        tools = load_tools(["serpapi"], llm=llm, serpapi_api_key=serpapi_api_key)
        questions = f"""我需要你为我编写1道关于 {knowledge}知识的问题, 
                        参考现有的题目{example_problem}，你可以利用搜索引擎检查问题是否符合生活实际
                        难度为 {difficulty}，{otherRestriction}，只需要给出问题即可,以问句结尾"""
        # 定义prompt
        template = ChatPromptTemplate.from_messages([
        ("system", "你是一名数学教研专家，下面的问题请你用中文回答"),
        ("human", "{user_input}"),
        MessagesPlaceholder("agent_scratchpad")
        ])
        agent = create_openai_tools_agent(llm, tools, template)
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
        generated_problem = agent_executor.invoke({"user_input": questions})

        print("生成的题目: \n", generated_problem)
        question_list.append(generated_problem["output"])

    return question_list


def solve_problems(generated_questions):
    # First, let's load the language model we're going to use to control the agent.
    llm = ChatOpenAI(temperature=0.5, model_name="gpt-4-turbo-preview")

    # Next, let's load some tools to use. Note that the `llm-math` tool uses an LLM, so we need to pass that in.
    tools = load_tools(["serpapi", "llm-math"], llm=llm)

    # Finally, let's initialize an agent with the tools, the language model, and the type of agent we want to use.
    agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)

    # 初始化存储问题和答案的列表
    question_with_answer = []
    for question in generated_questions["question_contents"]:
        solution = agent.run(
            question + "返回解题过程和答案，计算时只能提供单行的python数学表达式，这样子的表达式无法被计算'470 // 60, 470 % 60'，如果有多个要计算的表达式请你逐步计算")
        question_with_answer.append({"content": question, "description": solution})
    print(question_with_answer)
    return question_with_answer


class QuestionFormat(BaseModel):
    content: str = Field(description="问题内容")
    options: Dict[str, str] = Field(description="题目选项,键为ABCD,值为选项内容")
    description: str = Field(description="题目解析")
    correct_answer: str = Field(description="正确答案选项")


def add_choices(question_with_answer):
    model = ChatOpenAI(temperature=0.5, model_name="gpt-4-turbo-preview")

    # Set up a parser + inject instructions into the prompt template.
    parser = JsonOutputParser(pydantic_object=QuestionFormat)

    template = """现在为你提供针对小初学生的数学问题和答案，请你把其改造为选择题。问题为：{question},答案为：{answer},
                {format_instructions}
                """
    prompt = PromptTemplate(
        template=template,
        input_variables=["question", "answer"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )

    chain = prompt | model | parser
    # 初始化题目数组
    question_completed = []
    for item in question_with_answer:
        ans = chain.invoke({"question": item["content"], "answer": item["description"]})
        ans["selected"] = -1
        ans["like"] = 0
        ans["edit"] = 0
        question_completed.append(ans)
        print(ans)

    return question_completed


def knowledge_format(knowledge):
    last_comma_index = knowledge.rfind(',')

    # Check if there is a comma in the string
    if last_comma_index != -1:
        # Get the content after the last comma
        content_after_last_comma = knowledge[last_comma_index + 1:]
    else:
        # If there is no comma, the whole string is the content
        content_after_last_comma = knowledge

    return content_after_last_comma


@problem_bp.errorhandler(400)
def handle_bad_request(e):
    return jsonify(error=str(e)), 400


@problem_bp.route('store', methods=["POST"])
def problem_store():
    user_id = request.form.get('user_id')
    if user_id is None:
        return jsonify({'error': 'No user_id provided'}), 400
    export_list = request.form.getlist('exportList[]')
    print(export_list)
    hour = request.form.get('hour')
    minute = request.form.get('minute')
    second = request.form.get('second')
    name = request.form.get('name')
    bank = Bank(bank_user=user_id, bank_hour=hour, bank_minute=minute, bank_second=second, bank_name=name)
    db.session.add(bank)
    db.session.flush()
    bank_id = bank.bank_id
    for item in export_list:
        question_bank = QuestionBank(bank_id=bank_id, question_id=item)
        db.session.add(question_bank)
    # 提交更改
    db.session.commit()
    return jsonify({'bank_id': bank_id}), 200


@problem_bp.route('edit', methods=["POST"])
def problem_edit():
    question_id = request.form.get('question_id')
    question_content = request.form.get('question_content')
    option_list = request.form.getlist('optionList[]')
    correct_answer = request.form.get('correct_answer')
    description = request.form.get('description')

    session = db.session

    # 改questions表
    session.query(Question).filter_by(question_id=question_id).update({
        Question.question_text: question_content,
        Question.description: description
    })

    choices = db.session.query(Choice).filter_by(question_id=question_id).all()
    i = 0
    option_array = ['A', 'B', 'C']
    for choice in choices:
        if correct_answer == option_array[i]:
            choice.is_correct = 1
        else:
            choice.is_correct = 0

        choice.choice_text = option_list[i]
        i += 1

    # 改choices表
    for new_choice in choices:
        session.query(Choice).filter_by(choice_id=new_choice.choice_id).update({
            Choice.choice_text: new_choice.choice_text,
            Choice.is_correct: new_choice.is_correct
        })

    # 提交会话以保存更改
    session.commit()

    # 关闭会话
    session.close()
    return jsonify({'msg': 'edit success'}), 200


@problem_bp.route('/answer', methods=["POST"])
def answer_problem():
    try:
        question_list = request.form.getlist('question_array[]')
        answer_list = request.form.getlist('answer_array[]')
        leftTime = request.form.get('leftTime')  # Expected in the format "HH:MM:SS"
        bank_id = request.form.get('bank_id')
        # Convert leftTime to a datetime.time object
        hours, minutes, seconds = map(int, leftTime.split(':'))
        exercise_time = time(hours, minutes, seconds)

        # Current time for finish datetime
        current_time = datetime.now()

        # Create a new Exercise record
        new_exercise = Exercise(
            exercise_finish_datetime=current_time,
            exercise_cost_time=exercise_time,
            exercise_bank=bank_id
        )
        db.session.add(new_exercise)
        db.session.flush()  # Flush to get the exercise_id for the new record

        # Record each answer
        for question_id, answer in zip(question_list, answer_list):
            answer_id = int(answer) if answer.isdigit() else 0
            new_answer = Answer(question_id=question_id, choice=answer_id, exercise_id=new_exercise.exercise_id)
            db.session.add(new_answer)

        db.session.commit()
        return jsonify({'status': 'success', 'message': 'Exercise and answers saved successfully.',
                        'exercise_id': new_exercise.exercise_id}), 200

    except Exception as e:  # Catching a more general exception
        db.session.rollback()
        return jsonify({'status': 'error', 'message': str(e)}), 500
    return jsonify("success"), 200


@problem_bp.route('getProblems', methods=["POST"])
def get_problems():
    bank_id = request.form.get('bank_id')
    print(bank_id)
    questions_query = (
        db.session.query(Question)
        .join(QuestionBank, QuestionBank.question_id == Question.question_id)
        .filter(QuestionBank.bank_id == bank_id)
        .all()
    )
    print(questions_query)
    # Format the result into the desired structure
    formatted_questions = []
    for question in questions_query:
        formatted_question = {
            "content": question.question_text,
            "correct_answer": None,
            "id": str(question.question_id),
            "options": {},
            "selected": 0,
            "description": question.description
        }

        # Query choices for the current question
        choices_query = db.session.query(Choice).filter(Choice.question_id == question.question_id).all()

        # Format choices into options
        option_letter = 'A'
        for choice in choices_query:
            if choice.is_correct:
                formatted_question["correct_answer"] = option_letter

            formatted_question["options"][option_letter] = {'text': choice.choice_text, 'id': choice.choice_id}

            option_letter = chr(ord(option_letter) + 1)  # Increment letter to the next one

        formatted_questions.append(formatted_question)
    print(formatted_questions)
    bank_query = db.session.query(Bank).filter(Bank.bank_id == bank_id).first()
    print(bank_query)
    return jsonify({"questions": formatted_questions,
                    "hour": bank_query.bank_hour,
                    "minute": bank_query.bank_minute,
                    "second": bank_query.bank_second
                    })


def store_question(parsed_problems, knowledge, difficulty):
    session = db.session
    for problem in parsed_problems:
        # 创建 Question 实例
        new_question = Question(
            question_text=problem['content'],
            question_type="multiple-choice",
            difficulty_level=difficulty,
            tag=knowledge,
            subject="Math",
            description=problem['description']

            # 设置其他必需的字段，如 question_type, difficulty_level 等
        )
        session.add(new_question)
        session.flush()  # 刷新以获取 question_id
        problem['id'] = new_question.question_id
        # 创建 Choice 实例
        for choice_text, choice in problem['options'].items():
            new_choice = Choice(
                question_id=new_question.question_id,
                choice_text=choice,
                is_correct=1 if choice_text == problem['correct_answer'] else 0
            )
            session.add(new_choice)

        # 如果需要，也可以创建 Answer 实例

    # 提交会话以保存更改
    session.commit()

    # 关闭会话
    session.close()
    return parsed_problems

@problem_bp.route('/getBankDetails/<int:bank_id>', methods=["GET"])
def get_bank_details(bank_id):
    try:
        bank = db.session.query(Bank).filter_by(bank_id=bank_id).first()
        if not bank:
            return jsonify({'error': 'Bank not found'}), 404

        questions_query = (
            db.session.query(Question)
            .join(QuestionBank, QuestionBank.question_id == Question.question_id)
            .filter(QuestionBank.bank_id == bank_id)
            .all()
        )

        formatted_questions = []
        for question in questions_query:
            choices_query = db.session.query(Choice).filter_by(question_id=question.question_id).all()
            options = {chr(65 + i): choice.choice_text for i, choice in enumerate(choices_query)}
            correct_answer = next((chr(65 + i) for i, choice in enumerate(choices_query) if choice.is_correct), None)

            formatted_questions.append({
                'question_id': question.question_id,
                'content': question.question_text,
                'difficulty_level': question.difficulty_level,
                'options': options,
                'correct_answer': correct_answer,
                'selected': 0,  # 或其他逻辑来确定是否选中
                'description': question.description
            })

        return jsonify({
            'selected_questions': formatted_questions
        }), 200

    except Exception as e:
        db.session.rollback()
        traceback.print_exc()  # 打印完整的堆栈跟踪
        return jsonify({'error': 'Internal Server Error'}), 500


def download_docx(download_list):
    doc = Document()
    doc.add_heading('练习题', 0)

    # 遍历问题列表
    for question_id in download_list:
        # 查询问题
        question = db.session.query(Question).filter(Question.question_id == question_id).first()
        if question:
            doc.add_heading(f"题目：{question.question_text}", level=1)
            # 查询当前问题的选项
            choices = db.session.query(Choice).filter(Choice.question_id == question.question_id).all()
            option_letter = 'A'
            for choice in choices:
                # 添加选项到文档
                doc.add_paragraph(f" {option_letter}：{choice.choice_text}")
                option_letter = chr(ord(option_letter) + 1)  # Increment letter to the next one

    fake_file = io.BytesIO()
    doc.save(fake_file)
    fake_file.seek(0)

    return send_file(fake_file, mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document')


def download_excel(download_excel):
    pass


@problem_bp.errorhandler(500)
def handle_internal_server_error(e):
    print(e)
    return jsonify(error=str(e), stack_trace=traceback.format_exc()), 500
